Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add constraints and runtime API to transpose #2322

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

arminaleTT
Copy link
Contributor

@arminaleTT arminaleTT commented Feb 27, 2025

Ticket

#2314

Problem description

The optimizer needs more ops with constraints and runtime support to be able to ingest real models. Transpose is useful for both rennet (#2277) and llama (#2084)

What's changed

Added constraints and runtime APIs to the transposeOp. Added unit tests to exercise the new API

Closes #2314

Checklist

  • New/Existing tests provide coverage for changes

Copy link
Contributor

@odjuricicTT odjuricicTT left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great! Thanks!

CreateTiledLayout(tensorShape, mlir::tt::ttnn::BufferType::L1,
mlir::tt::ttnn::TensorMemoryLayout::Interleaved);
auto legalExp = Device::getDeviceConstraints(workerGrid);
EXPECT_TRUE(static_cast<bool>(legalExp));
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add one testcase where the provided layout is invalid?

Copy link
Contributor

@azecevicTT azecevicTT left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Dialect changes looks good, few minor comments.

assert(inputs.size() == 1);

const auto inputShape =
mlir::cast<RankedTensorType>(getOperand().getType()).getShape();
Copy link
Contributor

@azecevicTT azecevicTT Feb 28, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can use getInput().getType() (you can check the name of the operands/attributes in TTNNOps.td, for every operand you have an appropriately generated getter) and you will directly get the RankedTensorType, getOperand() is reserved for more generic uses where you don't know exactly the operation you are working on.

assert(inputs.size() == 1);

const auto inputShape =
mlir::cast<RankedTensorType>(getOperand().getType()).getShape();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Add constraints and runtime API for transpose
3 participants